# -*- coding: utf-8 -*-

import torch
from transformers import BertTokenizer, BertModel
import numpy as np

# 加载 BERT 模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese')

def extract_keywords(text, top_n=5):
    """使用 BERT 提取关键词"""
    # 分词和编码
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)

    # 获取 BERT 特征
    with torch.no_grad():
        outputs = model(**inputs)

    # 取最后一层隐藏状态，确保正确维度
    last_hidden_states = outputs.last_hidden_state  # 形状为 (batch_size, sequence_length, hidden_size)

    # 如果 batch_size 是 1，使用 squeeze() 进行维度处理
    if last_hidden_states.dim() == 3:  # 确保是 3D 张量
        last_hidden_states = last_hidden_states.squeeze(0)  # 变为 (sequence_length, hidden_size)

    # 计算每个 token 的得分（这里使用平均池化）
    token_scores = torch.mean(last_hidden_states, dim=1).numpy()  # 变为 (sequence_length,)

    # 计算每个 token 的重要性得分（如 L2 范数）
    token_importance = np.linalg.norm(token_scores, axis=1)

    # 获取 tokens 和对应的得分
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0].numpy())

    # 创建词与得分的字典
    token_score_dict = {token: score for token, score in zip(tokens, token_importance)}

    # 去除特殊 token
    token_score_dict = {k: v for k, v in token_score_dict.items() if k not in ['[CLS]', '[SEP]', '[PAD]', '[UNK]']}

    # 根据得分排序并选出前 N 个关键词
    sorted_tokens = sorted(token_score_dict.items(), key=lambda item: item[1], reverse=True)
    keywords = [token for token, score in sorted_tokens[:top_n]]

    return keywords

# 示例文本
text = "机器学习是人工智能的一个分支。自然语言处理是机器学习的重要应用。"

# 提取关键词
keywords = extract_keywords(text)
print("提取的关键词：", keywords)

